import Image
import numpy as np
import JT
import GDMetric
import MIMetric
from scipy.ndimage import *


class Registration():
    """
    Usage:
    either call: SetFixedFileName(pathToImage) or SetFixedImage(itkImage)
    then call either: SetMovingFileName(pathToImage) or SetMovingImage(itkImage)
    a call to GetMetric(), will return the metric value.

    The components are initialized with each call to GetMetric, so no registration
    information is saved between calls to this class.

    Can be used to write image to file and read in to the image registration
    process each time instead of trying to pass the image in memory.

    """

    def __init__(self):
        """
        """
        #specify amount of blur for images.
        self._blur = 0
        
        # Setup default image metric
        self._imageMetric = GDMetric
        self._metricName = 'GradientDifference'
        
        # Setup a default optimizer
        self._optimizer = JT.Optimizer()    


    def SetFixedFileName(self,filename):
        """
        Read a fixed image using Python image library
        """
        im = Image.open(filename)
        self._fixedImage = im.convert('L')

    def SetMovingFileName(self,filename):
        """
        Read a moving image using Python image library
        """
        im = Image.open(filename)
        self._movingImage = im.convert('L')

        
    def SetImageMetric(self,metric):
        """
        Set the image metric to use, should have a method to 'GetValue'
        that takes two images (numpy arrays) and returns a scalar 
        value for the optimizer.
        Eg. value = imageMetric.GetValue(image1,image2)
        """
        self._imageMetric = metric
        
    
    def SetOptimizer(self,optimizer):
        """
        Connect the optimizer to the registration class.  This interface is 
        in flux, beware of changes.
        
        The current implementation is for the user to setup all registration 
        components (metric, optimizer, images) in the registration class 
        and then call 'Start' (in the registration class) to initiate the 
        registration.  This will in turn call the optimize function.
        """
        self._optimizer = optimizer

    
    def Start(self):
        """
        """
        pass
        

    def GetValue(self,pose):
        """
        return the current value of the Gradient Difference image metric.  Set the position and
        orientation of the transform.  Pose is a vector made up of
        (X translation, Y translation, Z rotation in degrees)
        """
        image1 = self._fixedImage
        
        # First the image is shifted, then rotated (about center)
        # For ndimage shift translation is read in as Y,X ???
        image2 = shift(self._movingImage,(pose[1],pose[0]))
        image2 = rotate(image2,pose[2],reshape=False)
        
        # Blur may be useful for smooting out noisy image metrics
        if self._blur:
            image1 = gaussian_filter(image1, self._blur)
            image2 = gaussian_filter(image2, self._blur)
        
        # Calculate the metric    
        try:
            val=self._imageMetric.GetValue(image1,image2)
        except:
            print "Error calculating metric"
            raise

        return val


